from system.fsm import Fsm
from system.eventServer import eventServer
from system.dynload import dynload

import random
from copy import copy

class Ga(Fsm):
    def __init__(self, parent, type = None, name = None):
        Fsm.__init__(self, parent, ['startEvolution','returnFitness'], type, name)
        
        self.defaultVar("numGens",0)
        self.defaultVar("popSize",100)
        self.defaultVar("genomeSize",100)
        
        self.defaultVar("generate","ai.bin.simpleGenerate")
        self.defaultVar("crossover","ai.bin.simpleCrossover")
        self.defaultVar("mutate","ai.bin.simpleMutate")
        self.defaultVar("select","ai.bin.tournamentSelect")
        self.defaultVar("combine","ai.bin.simpleCombine")
        
        self.defaultVar("crossoverProb",0.75)
        self.defaultVar("mutateProb",0.01)
        
        self.defaultVar("distEval",False)
        self.defaultVar("seed",None)
        
    def startEvolution_idle(self, cmd, args):
        if args.has_key('path') and args['path'] != self.getVar("path"):
            return
        
        copyVars = ['numGens','popSize','genomeSize','crossover','crossoverProb','mutate','mutateProb','combine','generate','select']
        
        for var in copyVars:
            if args.has_key(var):
                self.setVar(var,args[var])
                
        self.random = random.Random(self.getVar("seed"))
                
        self.generate = dynload(self.getVar("generate"))
        self.select = dynload(self.getVar("select"))
        self.crossover = dynload(self.getVar("crossover"))
        self.mutate = dynload(self.getVar("mutate"))
        self.combine = dynload(self.getVar("combine"))
        
        assert self.generate and self.select and self.crossover and self.mutate and self.combine
        
        self.population = [self.generate(self.getVar("genomeSize"),self.random) for j in range(self.getVar("popSize"))]
        
        self.gen = 0
        self.bestInd = None
        
        self.startGeneration()
        
    def startGeneration(self):
        for i, ind in enumerate(self.population):
            if self.getVar("distEval"):
                pull = eventServer.distPull
            else:
                pull = eventServer.pull
                
            pull("evaluateIndividual",{'id':i,'ind':ind,'gen':self.gen})
            
        self.fitness = {}
        self.transition("evaluating")
        
    def finishGeneration(self):
        for i, ind in enumerate(self.population):
            fit = self.fitness[i]
            if not self.bestInd or fit > self.bestInd[1]:
                self.bestInd = copy(ind), fit, self.gen
                #print self.bestInd
                #print "best", self.gen, fit
                
        self.population = self.combine(self.population,self.fitness,self.select,self.crossover,self.getVar("crossoverProb"),self.mutate,self.getVar("mutateProb"),self.random)
        
        self.gen += 1
        if self.getVar("numGens") <= 0 or self.gen < self.getVar("numGens"):
            self.startGeneration()
        else:
            print "ga done", self.bestInd
        
    def returnFitness_evaluating(self, cmd, args):
        self.fitness[args['id']] = args['fitness']
        
        if len(self.fitness) >= len(self.population):
            self.finishGeneration()
